'''
Created on Aug 10, 2009
@author: mjbommar
'''

import pylab, numpy, scipy, os
from tfaces.analysis.TimeSeriesDatabase import TimeSeriesDatabase

class Plot(TimeSeriesDatabase):
    '''
    The Plot class extends the Database object to provide visualization
    from the import underlying data source.
    '''
    
    def __init__(self, dbFile, independentVariable = None):
        '''
        Constructor
        '''
        TimeSeriesDatabase.__init__(self, dbFile, independentVariable)
        
        self.figurePath = '%s_figures' % dbFile
        
        if not os.path.exists(self.figurePath): 
            os.mkdir(self.figurePath)

    def plotRunBySeed(self, seed, outputs = None):
        '''
        Plot the outputs of a run specified by its seed.
        '''
        M = self.getRunBySeed(seed, outputs)
    
        # If the outputs were passed, then we need to get the time index and label.
        if outputs:
            if self.independentVariable not in outputs:
                raise Exception("Plot::plotRunBySeed: independent variable not included")
            
            tmpIndex = outputs.index(self.independentVariable)
            timeLabel = outputs[tmpIndex]
        else:
            outputs = self.outputs
            tmpIndex = self.independentVariableIndex
            timeLabel = outputs[self.independentVariableIndex]

        # Set up the data index and labels
        dataIndex = [i for i in range(len(outputs)) if i != tmpIndex]
        dataLabels = [outputs[i] for i in dataIndex]
        
        # Plot the data now.
        pylab.plot(M[:, tmpIndex], M[:, dataIndex])
        pylab.xlabel(timeLabel)
        pylab.legend(dataLabels)
        pylab.show()
    
    def plotAverageRunByParameterSetID(self, id, outputs = None, errorbars = 0):
        '''
        Plot the average of the runs of a given parameter set.
        If the list of outputs is passed, then only these outputs will be plotted.
        If errorbars is greater than 0, then error bars of errorbars * stdev will be shown for each tick.
        '''
        M = self.getAverageRunByParameterSetID(id, outputs)
        
        if errorbars > 0:
            S = self.getStdRunByParameterSetID(id, outputs)
    
        # If the outputs were passed, then we need to get the time index and label.
        if outputs:
            if self.independentVariable not in outputs:
                raise Exception("Plot::plotRunBySeed: independent variable not included")
            
            tmpIndex = outputs.index(self.independentVariable)
            timeLabel = outputs[tmpIndex]
        else:
            outputs = self.outputs
            tmpIndex = self.independentVariableIndex
            timeLabel = outputs[self.independentVariableIndex]

        # Set up the data index and labels
        dataIndex = [i for i in range(len(outputs)) if i != tmpIndex]
        dataLabels = [outputs[i] for i in dataIndex]
        
        # Plot the data now.
        if errorbars > 0:
            for col in dataIndex:
                pylab.errorbar(M[:, tmpIndex], M[:, col], errorbars * S[:, col])
        else:
            pylab.plot(M[:, tmpIndex], M[:, dataIndex])
        
        pylab.xlabel(timeLabel)
        pylab.legend(dataLabels)
        pylab.show()



'''
p = Plot('rlr-2.db')
#p.plotRunBySeed(959077218, ['time','AvgDisttoSource'])
p.plotAverageRunByParameterSetID(1, ['time', 'AvgDisttoSource'], 1.95)


d = Database('rlr-2.db')
for id in d.getParameterSetByParameters({'stopT':6000}):
    M = d.getAverageRunByParameterSetID(id)
    pylab.plot(M[:, d.independentVariableIndex], M[:,1])
    pylab.show()
    break
'''